import numpy as np
import tensorflow as tf # 假设服务器和客户端使用TensorFlow模型
import copy
import time # 可选：用于记录优化时间

# --- 辅助函数 ---

def flatten_weights(weights_list):
    """将Keras模型的权重列表（通常是多维NumPy数组）展平成一个单一的一维NumPy数组。"""
    return np.concatenate([w.flatten() for w in weights_list])

def unflatten_weights(flat_weights, shapes_and_dtypes):
    """将扁平化的一维NumPy数组恢复成原始的Keras模型权重列表结构。"""
    unflattened = []
    current_pos = 0
    for shape, dtype in shapes_and_dtypes:
        # 检查形状是否有效
        if np.any(np.array(shape) <= 0):
            raise ValueError(f"无效的权重形状: {shape}")
        size = np.prod(shape)
        # 检查是否有足够的元素来重塑
        if current_pos + size > len(flat_weights):
            raise ValueError(f"扁平数组长度不足以填充形状 {shape}."
                             f"需要 {size} 元素，但只有 {len(flat_weights) - current_pos} 剩余。")

        layer_weights = flat_weights[current_pos:current_pos + size].reshape(shape).astype(dtype)
        unflattened.append(layer_weights)
        current_pos += size
     # 检查是否所有元素都被使用
    if current_pos != len(flat_weights):
        # 这个警告可能表示 shapes_and_dtypes 与 flat_weights 不匹配
        print(f"[Warning] 扁平数组在解压后仍有 {len(flat_weights) - current_pos} 个剩余元素。")
        # raise ValueError("扁平数组的长度与目标形状不完全匹配。")

    return unflattened


def get_shapes_and_dtypes(weights_list):
    """获取模型权重的形状和数据类型列表，用于后续的解压（unflatten）。"""
    return [(w.shape, w.dtype) for w in weights_list]

def relu(x):
    """ReLU激活函数。"""
    return np.maximum(0, x)

# --- 核心攻击逻辑 ---

def calculate_fltrust_aggregate_flat(g0_flat, client_deltas_flat, g0_norm):
    """
    使用扁平化的权重计算FLTrust聚合后的更新。
    模拟FLTrust聚合规则（公式4/7），但操作在一维向量上。

    Args:
        g0_flat (np.ndarray): 扁平化的服务器更新向量 g0。
        client_deltas_flat (list[np.ndarray]): 扁平化的客户端更新向量列表 [g1, g2, ...]。
        g0_norm (float): 预先计算好的服务器更新g0的L2范数。

    Returns:
        np.ndarray: 扁平化的聚合后梯度向量 g_agg。
    """
    # 如果没有客户端更新，返回零向量
    if not client_deltas_flat:
        return np.zeros_like(g0_flat)

    total_TS = 0.0 # 初始化总信任分数
    weighted_normalized_deltas = [] # 存储加权并归一化后的delta

    # 为避免除零，给范数加上一个很小的正数 epsilon
    g0_norm_eps = g0_norm + 1e-9

    # 遍历每个客户端的更新
    for gi_flat in client_deltas_flat:
        gi_norm = np.linalg.norm(gi_flat)
        gi_norm_eps = gi_norm + 1e-9

        # 1. 计算余弦相似度: cos(theta_i) = <gi, g0> / (||gi|| * ||g0||)
        # 使用加了epsilon的范数
        cosine_sim = np.dot(gi_flat, g0_flat) / (gi_norm_eps * g0_norm_eps)

        # 2. 计算信任分数: TS_i = ReLU(cos(theta_i))
        TS_i = relu(cosine_sim)
        total_TS += TS_i

        # 3. 计算归一化后的更新 g_bar_i = (||g0|| / ||gi||) * gi
        #    并乘以信任分数 TS_i
        #    这里直接计算加权项: TS_i * g_bar_i = TS_i * (g0_norm / gi_norm_eps) * gi_flat
        weighted_normalized_deltas.append(TS_i * (g0_norm / gi_norm_eps) * gi_flat)

    # 4. 聚合: g_agg = Sum(TS_i * g_bar_i) / Sum(TS_j)
    aggregated_flat = np.zeros_like(g0_flat)
    # 同样为总信任分数加epsilon避免除零
    total_TS_eps = total_TS + 1e-9
    for weighted_delta in weighted_normalized_deltas:
        aggregated_flat += weighted_delta

    # 除以总信任分数完成加权平均
    return aggregated_flat / total_TS_eps

def compute_objective_h(e_prime_list_flat, # 列表，包含 m 个恶意方向向量 e'_1 到 e'_m (numpy数组)
                        e0_flat,            # 服务器更新的方向向量 e_0 (numpy数组)
                        benign_e_list_flat, # 列表，包含 n-m 个良性方向向量 e_{m+1} 到 e_n
                        benign_c_list_relu, # 列表，包含 n-m 个良性方向与e0的 *ReLU处理后* 的余弦相似度
                        g0_norm,            # 服务器更新 g0 的范数 ||g0||
                        s_vector_flat):     # 攻击前的良性聚合结果 g_benign 的符号向量 sgn(g_benign)
    """
    计算FLTrust论文中公式(10)定义的目标函数 h。
    目标是最大化 h，即使得恶意聚合结果 g' 与良性聚合结果 g 的差异尽可能大（沿着 -s 方向）。
    所有输入的向量都应是扁平化的一维NumPy数组。
    """
    m = len(e_prime_list_flat)          # 恶意客户端数量
    num_benign = len(benign_e_list_flat) # 良性客户端数量

    # --- 计算 g' (恶意聚合结果) 中的各项 ---
    # g' = ||g0|| * [ Sum_mal(ReLU(<e'_i, e0>) * e'_i) + Sum_benign(ReLU(<e_i, e0>) * e_i) ] / Denom_g'
    # Denom_g' = Sum_mal(ReLU(<e'_i, e0>)) + Sum_benign(ReLU(<e_i, e0>))

    # 计算恶意部分的信任分数之和 (分子和分母都需要)
    sum_relu_c_prime = 0.0
    malicious_contributions_to_g_prime = np.zeros_like(e0_flat)
    for i in range(m):
        ei_prime_flat = e_prime_list_flat[i]
        # <e'_i, e0>，假设 e' 和 e0 都是单位向量
        c_prime_i = np.dot(ei_prime_flat, e0_flat)
        relu_c_prime_i = relu(c_prime_i)
        sum_relu_c_prime += relu_c_prime_i
        malicious_contributions_to_g_prime += relu_c_prime_i * ei_prime_flat # 分子中的恶意项：TS'_i * e'_i

    # 良性部分的信任分数之和 (已经预先计算好并应用了ReLU)
    sum_relu_c_benign = np.sum(benign_c_list_relu)

    # g' 的分母
    denominator_g_prime = sum_relu_c_prime + sum_relu_c_benign + 1e-9 # 加epsilon防除零

    # 计算良性部分对 g' 的贡献 (分子项)
    benign_contributions_to_g_prime = np.zeros_like(e0_flat)
    for i in range(num_benign):
        ei_benign_flat = benign_e_list_flat[i]
        c_i_relu = benign_c_list_relu[i] # 已经应用ReLU
        benign_contributions_to_g_prime += c_i_relu * ei_benign_flat # 分子中的良性项：TS_i * e_i

    # 组合得到 g' (方向向量，未乘以 ||g0||)
    g_prime_direction = (malicious_contributions_to_g_prime + benign_contributions_to_g_prime) / denominator_g_prime

    # --- 计算 g (仅良性聚合结果) 中的项 ---
    # g = ||g0|| * Sum_benign(ReLU(<e_i, e0>) * e_i) / Denom_g
    # Denom_g = Sum_benign(ReLU(<e_i, e0>))
    denominator_g_benign = sum_relu_c_benign + 1e-9 # g 的分母只包含良性信任分数
    # g 的分子项就是上面计算的 benign_contributions_to_g_prime
    g_benign_direction = benign_contributions_to_g_prime / denominator_g_benign

    # --- 计算目标函数 h (公式10) ---
    # h = ||g0|| * s^T * (g_direction - g'_direction)
    # 注意论文公式是 s^T * [Term1 - Term2 - Term3]，这里 Term1=g, Term2+Term3=g'
    h_value = g0_norm * np.dot(s_vector_flat, g_benign_direction - g_prime_direction)
    return h_value
def craft_heuristic_malicious_updates(g0,                      # 服务器更新 (模型权重的列表)
                                      benign_deltas,         # 良性客户端更新的列表
                                      m,                       # 恶意客户端的数量
                                      shapes_dtypes,           # 用于解压权重的形状和类型信息
                                      target_cosine_sim): # 目标余弦相似度
    """
    使用启发式方法构造恶意更新，旨在规避ReLU过滤器并指向有害方向。

    Args:
        g0 (list): 服务器更新权重列表。
        benign_deltas (list): 良性客户端更新权重列表。
        m (int): 恶意客户端数量。
        shapes_dtypes (list): 权重形状和类型元组列表。
        target_cosine_sim (float): 目标恶意方向与服务器方向的余弦相似度 (应为正数)。

    Returns:
        list[list]: 包含 m 个最终恶意更新 (权重列表形式) 的列表。
    """
    print(f"[启发式攻击] 开始为 {m} 个恶意客户端构造更新 (目标余弦相似度: {target_cosine_sim})...")
    start_time = time.time()

    # 1. 准备：扁平化，计算 g0 范数和方向 e0
    g0_flat = flatten_weights(g0)
    g0_norm = np.linalg.norm(g0_flat)
    g0_norm_eps = g0_norm + 1e-9
    if g0_norm < 1e-9:
        print("[启发式攻击警告] 服务器更新 g0 范数接近于零，无法构造攻击。返回零更新。")
        zero_delta = [np.zeros(shape, dtype=dtype) for shape, dtype in shapes_dtypes]
        return [copy.deepcopy(zero_delta) for _ in range(m)]
    e0_flat = g0_flat / g0_norm_eps

    # 2. 计算良性聚合 g_benign (用于确定有害方向)
    benign_deltas_flat = [flatten_weights(delta) for delta in benign_deltas]
    # 如果没有良性客户端，无法计算g_benign，可以使用-g0作为有害方向
    if not benign_deltas_flat:
        print("[启发式攻击警告] 没有良性更新来确定有害方向，使用 -g0 作为替代。")
        g_benign_flat = -g0_flat # 备选策略
    else:
        g_benign_flat = calculate_fltrust_aggregate_flat(g0_flat, benign_deltas_flat, g0_norm)

    # 3. 确定目标有害方向 d_harm (良性聚合的反方向)
    d_harm_flat = -g_benign_flat
    # 或者使用更简单的 d_harm_flat = -g0_flat

    # 4. 计算 d_harm 与 e0 正交的分量
    # <d_harm, e0>
    dot_harm_e0 = np.dot(d_harm_flat, e0_flat)
    # d_harm 在 e0 上的投影
    d_harm_proj_e0 = dot_harm_e0 * e0_flat
    # 正交分量
    d_harm_ortho = d_harm_flat - d_harm_proj_e0
    d_harm_ortho_norm = np.linalg.norm(d_harm_ortho)

    # 处理正交分量为零的特殊情况 (d_harm 与 e0 平行)
    if d_harm_ortho_norm < 1e-9:
        print("[启发式攻击警告] 有害方向与服务器方向平行，生成随机正交方向。")
        # 生成一个随机向量，并使其与e0正交
        random_vec = np.random.randn(len(e0_flat))
        random_proj_e0 = np.dot(random_vec, e0_flat) * e0_flat
        d_harm_ortho = random_vec - random_proj_e0
        d_harm_ortho_norm = np.linalg.norm(d_harm_ortho)
        # 如果随机向量也恰好平行（极小概率），则无法进行
        if d_harm_ortho_norm < 1e-9:
             print("[启发式攻击错误] 无法生成有效的正交方向。返回零更新。")
             zero_delta = [np.zeros(shape, dtype=dtype) for shape, dtype in shapes_dtypes]
             return [copy.deepcopy(zero_delta) for _ in range(m)]

    # 归一化正交方向
    e_harm_ortho = d_harm_ortho / (d_harm_ortho_norm + 1e-9)

    # 5. 组合构造恶意方向 e'
    # 确保 target_cosine_sim 在合理范围 (0, 1] 内，这里简化为 (0, 1)
    c_target = np.clip(target_cosine_sim, 1e-6, 1.0 - 1e-6) # 限制在(0, 1)开区间内
    s_ortho_squared = 1 - c_target**2
    # 确保开方参数非负
    if s_ortho_squared < 0:
         s_ortho_squared = 0 # 理论上 c_target<=1 不会发生
         print("[启发式攻击警告] target_cosine_sim 计算根号参数小于零，已修正。")

    s_ortho = np.sqrt(s_ortho_squared)

    # 构造恶意方向 e' = s_ortho * e_ortho + c_target * e0
    malicious_e_prime_flat = s_ortho * e_harm_ortho + c_target * e0_flat

    # （可选）验证余弦相似度并重新归一化以防万一
    # actual_cos_sim = np.dot(malicious_e_prime_flat, e0_flat)
    # print(f"[启发式攻击信息] 构造出的方向与e0的余弦相似度: {actual_cos_sim:.4f}")
    malicious_e_prime_flat /= (np.linalg.norm(malicious_e_prime_flat) + 1e-9)

    # 6. 生成最终恶意更新 g' = ||g0|| * e'
    malicious_g_prime_flat = g0_norm * malicious_e_prime_flat
    malicious_g_prime_unflattened = unflatten_weights(malicious_g_prime_flat, shapes_dtypes)

    # 为所有 m 个恶意客户端生成相同的更新（简单启发式）
    final_malicious_updates = [copy.deepcopy(malicious_g_prime_unflattened) for _ in range(m)]

    end_time = time.time()
    print(f"[启发式攻击] 完成构造，耗时 {end_time - start_time:.2f} 秒。")
    return final_malicious_updates